{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Polychromatic Propagations\n", "\n", "Prysm has a long heritage solving the monochromatic problem very quickly. However, it used a brute force 'propagate and interpolate' approach to solving the polychromatic problem. v0.19 offers large speedup by using matrix triple product DFTs to perform polychromatic propagations. This results in forward model times that are significantly faster, and clearer code when propagating to a grid for a detector.\n", "\n", "This notebook shows in the fewest possible lines the speedup by using a zero OPD circular pupil under nine discrete wavelengths. It also shows propagating to a detector grid." ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "from operator import add\n", "from functools import reduce\n", "\n", "import numpy as np\n", "\n", "from matplotlib import pyplot as plt\n", "\n", "from prysm import Pupil, PSF\n", "from prysm.propagation import focus_units, Wavefront" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "wvls = np.linspace(.4, .7, 9) # 400 to 900 nm, 9 wavelengths\n", "pups = [Pupil(wavelength=w, samples=512) for w in wvls]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## The old way" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "2.32 s ± 121 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] } ], "source": [ "%%timeit\n", "psfs = [PSF.from_pupil(pup, efl=1, Q=2) for pup in pups]\n", "poly_psf = PSF.polychromatic(psfs)" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "# the new way - setup for equivalence of output\n", "x, y = focus_units(pups[0].fcn, pups[0].sample_spacing, 1, wvls[0], 2)\n", "sample_spacing = x[1] - x[0]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## The new way" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "458 ms ± 6.07 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] } ], "source": [ "%%timeit\n", "psf_fields = [p.astype(Wavefront)\\\n", " .focus_fixed_sampling(efl=1, sample_spacing=sample_spacing, samples=1024) for p in pups]\n", "psf_intensities = [abs(w.fcn)**2 for w in psf_fields]\n", "poly_psf2 = PSF(x=x, y=y, data=reduce(add, psf_intensities))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "In this simple example, the new way is about **4x** faster. In this case, the simulation was done at Q=2 for all colors in the 'old' polychromatic way. This results in some numerical errors, where the new way is error free. At larger Qs the old way will have improved accuracy, but also increased computation time. To show the true power of the new way, we will compare old and new for Q=8, and use the flexibility of the matrix triple product to compute a smaller output domain:" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## the old (high oversampling):" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "48.8 s ± 265 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] } ], "source": [ "%%timeit\n", "psfs = [PSF.from_pupil(pup, efl=1, Q=8) for pup in pups]\n", "poly_psf3 = PSF.polychromatic(psfs)" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "# the new way - setup for equivalence of output\n", "x, y = focus_units(pups[0].fcn, pups[0].sample_spacing, 1, wvls[0], 8)\n", "sample_spacing = x[1] - x[0]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## the new (high oversampling)" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "103 ms ± 4.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] } ], "source": [ "%%timeit\n", "psf_fields = [p.astype(Wavefront)\\\n", " .focus_fixed_sampling(efl=1, sample_spacing=sample_spacing, samples=256) for p in pups]\n", "psf_intensities = [abs(w.fcn)**2 for w in psf_fields]\n", "poly_psf4 = PSF(x=x, y=y, data=reduce(add, psf_intensities))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The parameters of this second example are not particularly relevant outside coronagraphy or simulation study of astronomical instrumentation due to the large Q, but we see a **nearly 500x** speedup for use of the matrix triple product.\n", "\n", "While the output data is not strictly the same since the matrix triple product is computed over a smaller domain, their _value_ is the same since we do not care about the region far from the core of the PSF. This speedup allows computations that may require a supercomputer to be done on a laptop." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.8.5" } }, "nbformat": 4, "nbformat_minor": 2 }